import numpy as np
import matplotlib.pyplot as plt
import tqdm
import os

if __name__ == "__main__":
    path = "/disk527/DataDisk/a804_cbf/datasets/lunar_crater/Lunar_LRO_LOLAKaguya_DEMmerge_60N60S_512ppd.tif"
    output_dir = "/disk527/DataDisk/a804_cbf/datasets/lunar_crater"
    LABEL_BYTES = 984041
    BYTES_PER_PIXEL = 1
    LINES = 61440
    LINE_SAMPLES = 184320
    base_batch = 960
    num_batches = LINES // base_batch
    num_samples = LINE_SAMPLES // base_batch

    print(
        f"The basic batch of DEM is {base_batch}x{base_batch}. Each line contains {num_samples} batches. There are totally {num_batches*num_samples} basic batches."
    )

    with open(path, "rb") as f:
        # 跳过前984003的标签字节
        line = f.read(LABEL_BYTES)

        for row in tqdm.tqdm(range(num_batches)):
            # 读取一个标准batch的数据，即960行，184320列，共计192个batch
            data = f.read(LINE_SAMPLES * BYTES_PER_PIXEL * base_batch)
            # 将数据转换为numpy数组
            data_array = np.frombuffer(data,dtype=np.uint8).reshape(
                base_batch, LINE_SAMPLES
            )

            if not os.path.exists(os.path.join(output_dir, "textures")):
                os.makedirs(os.path.join(output_dir, "textures"))
            if not os.path.exists(os.path.join(output_dir, "textures", f"{row}")):
                os.makedirs(os.path.join(output_dir, "textures", f"{row}"))

            for col in range(num_samples):
                img_array = data_array[:, col * base_batch : (col + 1) * base_batch]
                # 输出深度图以供查验
                plt.imsave(
                    os.path.join(
                        output_dir,
                        "textures",
                        f"{row}",
                        f"txr_{row}_{col}_{base_batch}.png",
                    ),
                    img_array,
                    cmap="gray",
                )
